from dsec.dataset import DSECDet
import numpy as np
import torch
import cv2


dataset = DSECDet(
    "E:/Datasets/DSEC",
    "test"
)
print(len(dataset))
image, ev_repr, label_bboxes, label_classes = dataset[2800]
print(image.shape)
img: np.ndarray = image.permute(1, 2, 0).numpy().astype(np.uint8)
cv2.imwrite("img1.png", img)

print(ev_repr.shape)
ev_repr = torch.sum(ev_repr, dim = 0).numpy().astype(np.uint8)
ev_img: np.ndarray = np.zeros_like(img)
for i in (0, -1):
    ev_img[..., i] = ev_repr[i]
cv2.imwrite("img2.png", ev_img*2)

mask = np.sum(ev_repr, axis = 0)
mask = np.stack([mask] * 3).transpose(1, 2, 0)

img1 = img.copy()
ev_img1 = ev_img.copy()
# res1 = img1 * (mask / np.max(mask))
res1 = img1 * (np.clip(mask, 0, 10)/10).astype(np.uint8)
res2 = (img1 / 3) + (res1 / 2 * 3)
cv2.imwrite("img3.png", res1)
cv2.imwrite("img4.png", res2)